package main

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"database/sql"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"log"

	_ "github.com/go-sql-driver/mysql"
)

func GetKeysFromDB(db *sql.DB) (string, string, error) {
	var privateKey, publicKey string
	err := db.QueryRow("SELECT private_key, public_key FROM sys_client where app_code = ?", "go-server").Scan(&privateKey, &publicKey)
	if err != nil {
		return "", "", err
	}
	return privateKey, publicKey, nil
}

func main() {
	dsn := "root:123456@tcp(127.0.0.1:3306)/go_admin?charset=utf8mb4&parseTime=True&loc=Local"
	db, err := sql.Open("mysql", dsn)
	if err != nil {
		log.Fatalf("Error connecting to database: %v", err)
	}
	defer db.Close()

	privateKeyPEM, publicKeyPEM, err := GetKeysFromDB(db)
	if err != nil {
		log.Fatalf("Error retrieving keys from database: %v", err)
	}
	//解析私钥和公钥
	privateKey, err := parsePrivateKey(privateKeyPEM)
	if err != nil {
		log.Fatalf("Failed to parse private key: %v", err)
	}
	publicKey, err := parsePublicKey(publicKeyPEM)
	if err != nil {
		log.Fatalf("Failed to parse public key: %v", err)
	}

	// 要加密的明文数据
	plaintext := []byte("123456")

	// 公钥加密
	ciphertext, err := rsaEncrypt(publicKey, plaintext)
	if err != nil {
		log.Fatalf("Failed to encrypt: %v", err)
	}

	ctext := base64.StdEncoding.EncodeToString(ciphertext)
	fmt.Printf("CiphertextStr: %s\n", ctext)
	decodeString, err := base64.StdEncoding.DecodeString(ctext)
	// 私钥解密
	decryptedText, err := rsaDecrypt(privateKey, decodeString)
	if err != nil {
		log.Fatalf("Failed to decrypt: %v", err)
	}
	fmt.Printf("Decrypted Text: %s\n", decryptedText)
}

// 解析私钥
func parsePrivateKey(pemString string) (*rsa.PrivateKey, error) {
	block, _ := pem.Decode([]byte(pemString))
	if block == nil {
		return nil, fmt.Errorf("failed to parse PEM block containing the key")
	}
	privateKey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return nil, fmt.Errorf("failed to parse private key: %v", err)
	}
	return privateKey, nil
}

// 解析公钥
func parsePublicKey(pemString string) (*rsa.PublicKey, error) {
	block, _ := pem.Decode([]byte(pemString))
	if block == nil {
		return nil, fmt.Errorf("failed to parse PEM block containing the key")
	}
	publicKey, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, fmt.Errorf("failed to parse public key: %v", err)
	}
	rsaPublicKey, ok := publicKey.(*rsa.PublicKey)
	if !ok {
		return nil, fmt.Errorf("public key is not of type *rsa.PublicKey")
	}
	return rsaPublicKey, nil
}

// RSA公钥加密
func rsaEncrypt(publicKey *rsa.PublicKey, plaintext []byte) ([]byte, error) {
	ciphertext, err := rsa.EncryptPKCS1v15(rand.Reader, publicKey, plaintext)
	if err != nil {
		return nil, fmt.Errorf("failed to encrypt: %v", err)
	}
	return ciphertext, nil
}

// RSA私钥解密
func rsaDecrypt(privateKey *rsa.PrivateKey, ciphertext []byte) ([]byte, error) {
	plaintext, err := rsa.DecryptPKCS1v15(rand.Reader, privateKey, ciphertext)
	if err != nil {
		return nil, fmt.Errorf("failed to decrypt: %v", err)
	}
	return plaintext, nil
}
